[OMNIML-6079] [20/21] Add the Kimi-K3 expert-parallel adapter - #2749

Claude review — [20/21] Kimi-K3 expert-parallel adapter (round 4)

Scope: full review (the trigger comment carried no scoping instructions). 9 changed files (+1653/-4), all opened, in modelopt/ → modelopt_recipes/ → tests/ order. For context I re-read utils/moe.py (register_torch_a2a, check_distributed_support, apply_distributed_support, torch_a2a_experts_forward), model_load_utils.py (resolve_distribution_plan, _decide_distribution, the three apply_distributed_support call sites, _make_expert_param_matcher), _onfly_decompress.install_onfly_decompress, and quantization/plugins/huggingface.py (_running_grouped_mm, _local_experts, normalize_experts_impl_on_the_fly).

Prior rounds

Rounds 1–3 raised 2 IMPORTANT and 14 SUGGESTION findings; round 3 approved. Spot-checking the ones with code to check against the current tip (9df995b1): the ckpt_path is None refusal (ep.py:593), the live sentinel guard (num_experts + 1, ep.py:215), the VAR_KEYWORD handling plus the explicit cache_position-only drop-list (ep.py:406-419), reshape in the block forward, requires_transformers="5.13" on the spec, the loud raise when no KDA kernel is found (ep.py:338), the loud import of the conversion API (ep.py:540), and the per-checkpoint remote-module discovery (_find_module_defining, replacing the process-global _PATCHED flag Codex flagged) are all in place. CodeRabbit's install_onfly_decompress finding is addressed too (ep.py:641 gates on quantization_config without swallowing validation errors).

New findings this round

CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 2

IMPORTANT

1. _patch_kda_beta_sigmoid's _wrapped can still pass use_beta_sigmoid_in_kernel=True through with the raw gate (ep.py:359-371). Three fall-throughs leave the flag set and beta unsquashed — beta_pos is None (unintrospectable signature, or fla-core renaming the parameter while the remote code passes it positionally), a non-tensor beta keyword with beta_pos >= len(args), and a non-tensor args[beta_pos]. Nothing pins fla-core: the docstring names 0.5.0 but no floor enforces it, and the spec's transformers floor does not cover it. In each case fla-core ignores the flag, the delta-rule recurrence diverges to inf, and the only symptom is a NaN amax after a full EP64 calibration. This is the same argument you accepted one round ago for the sibling "no kernels to wrap" case, applied one level in: the kernel was found, but the gate it was wrapped for was not. Fix sketched in the thread.

SUGGESTION

2. prepare's side effects are process-global, but resolve_distribution_plan returns the raw config when the plan turns out non-EP (model_load_utils.py:1595) — fused block patched and converters registered, yet that config is neither marked for the decompressor nor stripped of quantization_config, so the fused experts would load missing. The ep_size == 1 and dense == "replicate" refusal at :1446 does not cover dense == "fsdp" with ep_size == 1, nor the early auto_class return at :1368. Family-wide (kimi_k25 / kimi_vl share it), so not a blocker — but Kimi-K3 is the first adapter that patches the checkpoint's own class and leaves it in place.
3. hasattr(Cache, "get_query_offset") reads through the MRO, so an inherited (and, for this hybrid cache, non-working) implementation suppresses the new shim; Cache.__dict__ is the test _shim_tie_weights_kwargs already uses for tie_weights (ep.py:452-463).

What checked out

• Fused layout and the a2a contract. gate_up_proj [E, 2I, H_latent] / down_proj [E, H_latent, I] with is_transposed=False, has_gate=True, has_bias=False is what torch_a2a_experts_forward reads (moe.py:446-508); _apply_gate's chunk(2, dim=-1) split matches the converters' MergeModulelist(dim=0) + Concatenate(dim=1) row order (w1=gate first, w3=up second), and SituAndMul taking the full concat matches the per-expert reference.
• The new register_torch_a2a() in prepare (ep.py:630) lands before _experts_implementation is named and before the plan decision's meta build, which is the ordering the "torch_a2a is not supported" failure needed. It only runs on the EP path (model_load_utils.py:1571 gates prepare on world_size > 1 and _ep_degree(...) != 1), the registry writes are additive and idempotent, and load_and_prepare_ep registering again is a no-op. The utils/moe.py docstring edit is the whole code change there.
• Latent-MoE ordering in the rewritten _forward: down_proj → fused experts → optional norm → up_proj → reshape(orig_shape) → shared experts on the un-flattened identity, with use_latent_moe read off the block rather than a config field the checkpoint lacks. Matches the stand-in's per-expert reference, which the test compares against.
• get_mask_sizes / get_query_offset consistency: (ql + past_seen, 0) and past_seen are exactly what transformers' own DynamicCache returns, so the two agree on kv_offset.
• Conversion registration by model_type only (not class names), with the de-duplication reasoning for the nested kimi_linear VLM scopes; overwrite=True is safe for a non-native model type.
• Decompressor ordering and idempotence: converters registered before install_onfly_decompress patches get_model_conversion_mapping, quantization_config dropped after, and a re-prepare of the same checkpoint hits the installed is not None / compression is None early returns at _onfly_decompress.py:337-351 rather than re-installing.
• Preset exclusions: the five Kimi-K3 projection spellings and mm_projector are additive enable: false entries in the shared default unit, consistent with how gemma's and Llama-4's projector spellings already live there; recipes can still re-enable them. The new parametrized test asserts both the exclusions and that self_attn.o_proj stays quantized, so it would fail on an over-broad glob.
• Spec wiring: prepared_ep_plan=BASE_MODEL_EP_PLAN with no base_model_ep_plan/plan_config_attr mirrors kimi_k25, so apply_distributed_support attaches nothing itself and _spec_ep_plan still sees the plan — which is what makes the new EP-degree-1 refusal test pass.
• Tests: the stand-in follows the checkpoint's own shapes (stale A_log as [num_heads], get_mask_sizes wanting a tensor, tie_weights with the old signature, fla-core 0.5 kernels taking the flag through **kwargs), and the query-offset test covers both the add and the leave-alone branch. SituAndMul holding plain floats confirms the fused module gains no parameter the checkpoint has no key for.

Risk

Moderate, and well bounded. The adapter is additive, runs only on the expert-parallel path, and is backed by both unit tests and a real EP64 load whose export is byte-compared against the previous one. The one IMPORTANT is a defensive gap in a defensive patch; neither SUGGESTION blocks. Note the PR body's own ~620 lines figure is now ~646 for ep.py alone (~700 source across the three new files) — the self-contained-model-drop exemption in CONTRIBUTING covers it, but the number in the description is stale.

🤖 Generated with Claude Code

Claude review — [20/21] Kimi-K3 expert-parallel adapter (second pass)

Scope applied: full review (the trigger comment carried no scoping instructions). 9 changed files, 1674 additions — all reviewed: modelopt/torch/models/kimi_k3/{ep,specs,__init__}.py, modelopt/torch/models/__init__.py, modelopt/torch/utils/moe.py, modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml, and the three test files. Read for context: utils/moe.py (check_distributed_support, apply_distributed_support, _spec_ep_plan, register_torch_a2a, _SUB_CONFIG_ATTRS), _onfly_decompress.install_onfly_decompress, plugins/model_load_utils.py (the three apply_distributed_support call sites + resolve_distribution_plan), quantization/plugins/huggingface.py (normalize_experts_impl_on_the_fly), models/specs.py.

Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2 new (both non-blocking)

Still open from the previous pass (not re-filed, listed for tracking): the _init(self, config) signature vs. the siblings' *args, **kwargs; the third copy of the tie_weights kwargs shim; the auto_map lookup accepting only the two text-only auto classes; and the ordering comment at ep.py:643-649. That last one is worth a second look — it now states the registration must happen "BEFORE the decompressor patches get_model_conversion_mapping", but install_onfly_decompress installs a patched wrapper that calls orig(...) at load time, so it sees whatever is registered then and the order here is free. The code is correct either way; only the comment asserts a constraint that is not there.

What I traced this round

• The new register_torch_a2a() in prepare (the main delta). Correct and necessary. prepare names torch_a2a on text_config, and resolve_distribution_plan builds the model on meta from that config before load_and_prepare_ep registers anything — transformers rejects an unregistered experts implementation at build, which is the experts_implementation="torch_a2a" is not supported failure the description reports. The registration is idempotent via _TORCH_A2A_REGISTERED, adds only new names to ALL_EXPERTS_FUNCTIONS / ALL_PARALLEL_STYLES, and load_and_prepare_ep calls it again harmlessly. test_a_model_builds_from_the_prepared_config_before_the_expert_parallel_load reproduces the failure honestly — it clears _global_mapping and _TORCH_A2A_REGISTERED to put the process back in the state the plan decision finds, and uses a stock Qwen3-MoE because the check is transformers', not Kimi's.
• Setting _experts_implementation in prepare does not leak to non-EP paths. All three apply_distributed_support call sites are gated on route == _EXPERT_PARALLEL, and load_and_prepare_ep:2571 only writes the attribute onto configs that already have it — which is precisely why prepare must seed it. resolve_distribution_plan discards the prepared copy when the plan is not expert-parallel (return plan, prepared[0] if plan.is_expert_parallel else hf_config), and the EP-degree-1 whole-model case is refused outright by the new test_resolve_distribution_plan parametrization rather than silently loading with a stale torch_a2a.
• Plan placement. ep_spec leaves base_model_ep_plan unset and supplies prepared_ep_plan, so apply_distributed_support skips the declarative branch and plan_config_attr is correctly irrelevant; prepare puts the plan on text_config, which is one of the two spellings the EP load reads (_SUB_CONFIG_ATTRS = ("text_config", "llm_config")). _spec_ep_plan falls through to prepared_ep_plan, so a dense-whole run still resolves the routed experts — and the test asserts the attached plan equals the spec's, so the two cannot drift.
• Fused layout and the gate. gate_up_proj [E, 2I, H_latent] / down_proj [E, H_latent, I] with is_transposed=False matches what torch_a2a_experts_forward indexes. The converter chain MergeModulelist(dim=0) then Concatenate(dim=1) puts w1 in rows [0:I] and w3 in [I:2I], which is the half chunk(2, dim=-1) / SituAndMul treats as the gate; has_gate=True routes through _apply_gate rather than act_fn, so the custom situ SwiGLU survives the a2a path and not just the eager fallback. The dispatch test checks this against a real transformers backend, so a layout error shows up as an eager/dispatched mismatch rather than as silent numerics.
• Block forward equivalence, including identity being the unflattened input and reshape(orig_shape) landing after the latent up-projection, across the latent / plain / situ parametrizations. The eager fallback's one_hot(..., num_experts + 1) then permute(2,1,0) then where(mask[e]) indexing is consistent against [T, topk] weights, and the sentinel class is skipped rather than folded in.
• The KDA beta fix, the riskiest patch here, is right in substance: sigmoid in fp32 cast back to beta.dtype, the kernel flag cleared so fla-core cannot double-apply, beta found by keyword or by position, and both failure modes refused loudly — remote code without the kernels to wrap, and now a call whose beta the wrapper cannot locate. Finding 1 is a robustness nit on the position lookup, not a correctness break: every path still raises rather than running the recurrence on the raw gate.
• Decompressor wiring. install_onfly_decompress is now called only when a quantization_config is present, and a non-compressed-tensors one reaches its own ValueError rather than being swallowed. Re-preparing a config whose quantization_config was already dropped hits the compression is None early return; a fresh deepcopy with the config still present re-marks it. _compressedify keys off a .weight suffix under experts, so K3's w1/w3/w2 sources are rewritten, and the mark is cleared once converters are rewritten so a later save_pretrained does not inherit it.
• Validate before mutate. Both prepare refusals (ckpt_path is None, trust_remote_code=False) fire before any config or class is touched, and check_distributed_support enforces the 5.13 floor before prepare runs at all.
• Collective symmetry. Every patch, shim and refusal is rank-independent; resolve_distribution_plan wraps the spec application in dist.all_reduce("model spec", ..., all), so a rank that cannot apply it fails every rank instead of stranding the others in a collective.
• Recipe change. --qformat in examples/hf_ptq/hf_ptq.py reads presets.QUANT_CFG_CHOICES, which is built from these preset YAMLs — so the new exclusions really do apply to "every built-in format", not just --recipe runs, and the new preset test exercises exactly that entry point. The five projection names and mm_projector are distinctive enough not to collide with other families, and the test's o_proj assertion would catch an over-broad glob.
• No prompt-injection content in the diff.

Risk assessment — low-to-moderate, unchanged from the first pass and inherent rather than introduced: the adapter monkey-patches a third-party checkpoint's remote code, bounded by model_type == "kimi_k3" on the expert-parallel path, with one process-wide side effect (register_checkpoint_conversion_mapping("kimi_linear", ...)) scoped to a model type nothing else in the repo claims. The standing risk is remote-code drift, and this revision has moved further toward confronting it — per-revision module lookup by defined class, a hard refusal when the KDA kernels are not the ones it knows, and now a refusal when the beta gate cannot be located. The ~620-line ep.py exceeds the ~500-line budget, which I read as covered by the "self-contained drop such as a new model/backend that has no working intermediate state" carve-out; the description says so up front.

Approving: no blocking issues.

🤖 Generated with Claude Code